library(rjags)
library(runjags)
library(MASS)

###Input variables
load("vd_variables.RData")

###Dichotomous variables
vd<-vd_variables$vd
###Count variables
vu<-vd_variables$vu
###Continuous variables
vc<-vd_variables$vc
###Proportion variables
vp<-vd_variables$vp
###Ordinal variables
vo3<-vd_variables$vo3
vo4<-vd_variables$vo4
vo5<-vd_variables$vo5


####Create normalized initial values for latent construct as mean of normalized means
inits<-cbind(vd,vu,vc,vp,vo3,vo4,vo5)

m.i<-apply(inits,2,function(x){mean(na.omit(x))})
s.i<-apply(inits,2,function(x){sd(na.omit(x))})

for(j in 1:ncol(inits)){
    inits[,j]<-(inits[,j]-m.i[j])/s.i[j]
}

inits<-apply(inits,1,mean,na.rm=T)
xi<-rep(NA,length(inits))
xi[seq(1,length(xi),10)]<-inits[seq(1,length(xi),10)]

####initial values for ordinal thresholds
taustarO3<-t(apply(vo3,2,function(x){polr(as.factor(x) ~ inits,method="probit")$zeta}))
taustarO4<-t(polr(as.factor(vo4) ~ inits,method="probit")$zeta)
taustarO5<-t(apply(vo5,2,function(x){polr(as.factor(x) ~ inits,method="probit")$zeta}))



###list of model values
lmod<-list(yD=as.matrix(vd), yU=as.matrix(vu), yC=as.matrix(vc), yP=as.matrix(vp), yO5=as.matrix(vo5),
           yO4=vo4, yO3=as.matrix(vo3),b0=rep(0,2),
           B0=diag(1,2), BD0=diag(.1,2), tO5=rep(0,4), TO5=diag(1,4), tO4=rep(0,3), TO4=diag(1,3), 
           tO3=rep(0,2), TO3=diag(1,2), N=nrow(vd), ND=ncol(vd), NU=ncol(vu),
           NC=ncol(vc), NP=ncol(vp), NO5=ncol(vo5), NO3=ncol(vo3))


####8 chains
hmodel <- run.jags(method="parallel", model="idmodel.txt", inits=list(taustarO3=taustarO3, 
                                                                    taustarO4=c(taustarO4),
                                                                    taustarO5=taustarO5), 
                   monitor=c("xi", "betaD", "betaU", "betaC", "tauC", "betaP",
                             "tauP", "betaO5", "tauO5", "betaO4", "tauO4", "betaO3",
                             "tauO3"), data=lmod, n.chains=1, adapt=1000, burnin=10000, 
                   sample=500, thin=20, summarise=F,plots=F, modules=c("glm", "lecuyer"))

save(hmodel, file="id_draws.RData")


#########Convert to MCMC list
cds<-as.mcmc.list(hmodel) 

###Flip scale on chains for consistent directionality

cds[[2]]<- -cds[[2]]
cds[[4]]<- -cds[[4]]
cds[[8]]<- -cds[[8]]

###Convert to MCMC
cds<-as.mcmc(cds)

##Function to normalize draws by iteration
idbayes<-function(x){
  mx<-apply(x,1,mean)
  sdx<-apply(x,1,sd)
  idb<-matrix(nrow=nrow(x), ncol=ncol(x))
  for(j in 1:nrow(x)){
    idb[j,]<- (x[j,]-mx[j])/sdx[j]
  }
  return(idb)
}

###Normalize draws
xix<- idbayes(cds[,1:16472][[1]])

###Median
idm<-apply(xix,2,median)

###95% HPD intervals
idhpd<-HPDinterval(as.mcmc(xix))

####add country-year labels and rename variables
idlab<-vd_variables$idlab
idanalysis<-data.frame(idlab,idm,idhpd)
names(idanalysis)[6:8]<-c("Median.Estimate","Lower.95.HPD","Upper.95.HPD")

save(idanalysis,file="idanalysis.RData")
